import os
import asyncio
from lightrag.llm.openai import openai_embed
from lightrag.utils import EmbeddingFunc

async def debug_lightrag_embedding():
    """调试LightRAG的embedding函数调用"""
    
    print("=== 环境变量检查 ===")
    print(f"EMBEDDING_MODEL: {os.getenv('EMBEDDING_MODEL', 'NOT_SET')}")
    print(f"EMBEDDING_BINDING_HOST: {os.getenv('EMBEDDING_BINDING_HOST', 'NOT_SET')}")
    print(f"EMBEDDING_BINDING_API_KEY: {os.getenv('EMBEDDING_BINDING_API_KEY', 'NOT_SET')}")
    print(f"OPENAI_API_KEY: {os.getenv('OPENAI_API_KEY', 'NOT_SET')}")
    print(f"EMBEDDING_DIM: {os.getenv('EMBEDDING_DIM', 'NOT_SET')}")
    print(f"MAX_EMBED_TOKENS: {os.getenv('MAX_EMBED_TOKENS', 'NOT_SET')}")
    
    print("\n=== 测试参数 ===")
    model = os.getenv("EMBEDDING_MODEL", "BAAI/bge-m3")
    base_url = os.getenv("EMBEDDING_BINDING_HOST", "https://api.siliconflow.cn/v1/embeddings")
    api_key = os.getenv("EMBEDDING_BINDING_API_KEY") or os.getenv("OPENAI_API_KEY", "sk-xxjzwmnsxdodiqnfwtxkizjnplmexrpfomqzetjaeppvvrnl")
    
    print(f"使用的模型: {model}")
    print(f"使用的base_url: {base_url}")
    print(f"使用的api_key: {api_key[:20]}..." if api_key else "api_key: None")
    
    print("\n=== 直接测试openai_embed函数 ===")
    try:
        test_texts = ["This is a test string for embedding."]
        result = await openai_embed(
            texts=test_texts,
            model=model,
            base_url=base_url,
            api_key=api_key,
        )
        print(f"✅ 直接调用成功!")
        print(f"结果形状: {result.shape}")
        print(f"Embedding维度: {result.shape[1]}")
        
    except Exception as e:
        print(f"❌ 直接调用失败: {e}")
        import traceback
        traceback.print_exc()
    
    print("\n=== 测试EmbeddingFunc包装器 ===")
    try:
        embedding_func = EmbeddingFunc(
            embedding_dim=int(os.getenv("EMBEDDING_DIM", "1024")),
            max_token_size=int(os.getenv("MAX_EMBED_TOKENS", "8192")),
            func=lambda texts: openai_embed(
                texts,
                model=model,
                base_url=base_url,
                api_key=api_key,
            ),
        )
        
        test_texts = ["This is a test string for embedding."]
        result = await embedding_func(test_texts)
        print(f"✅ EmbeddingFunc包装器成功!")
        print(f"结果形状: {result.shape}")
        print(f"Embedding维度: {result.shape[1]}")
        
    except Exception as e:
        print(f"❌ EmbeddingFunc包装器失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    asyncio.run(debug_lightrag_embedding()) 